# ui/community_viewer.py

import os

from PyQt5.QtGui import QFont
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QListWidget
from matplotlib import pyplot as plt

from analysis.visualization import draw_network


class CommunityViewer(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout(self)

        title = QLabel("检测到的蛋白质社区", self)
        title.setFont(QFont('Arial', 24, QFont.Bold))
        layout.addWidget(title)

        # 社区列表
        self.community_list = QListWidget(self)
        self.load_communities()
        layout.addWidget(self.community_list)

        # 网络图展示区域
        self.figure_layout = QHBoxLayout()
        self.canvas_container = QWidget(self)
        self.canvas_container.setLayout(self.figure_layout)
        layout.addWidget(self.canvas_container)

        # 查看 GO 分析按钮
        self.btn_go_analysis = QPushButton("查看功能富集 (GO)", self)
        self.btn_go_analysis.clicked.connect(self.open_go_analysis)
        layout.addWidget(self.btn_go_analysis)

        # 返回主菜单按钮
        back_button = QPushButton("返回主菜单", self)
        back_button.clicked.connect(lambda: self.parent().switch_page(0))
        layout.addWidget(back_button)

        self.setLayout(layout)

    def load_communities(self):
        community_dir = "results/communities"
        if not os.path.exists(community_dir):
            self.community_list.addItem("未找到社区文件，请先运行 LAMA 社区检测")
            return

        for file in sorted(os.listdir(community_dir)):
            community_id = file.replace("lama_community_", "").replace(".txt", "")
            self.community_list.addItem(f"社区 {community_id}")

    def open_go_analysis(self):
        selected = self.community_list.currentItem()
        if selected:
            community_id = selected.text().replace("社区 ", "")
            self.parent().go_viewer.set_community_id(community_id)
            self.parent().switch_page(1)


class CommunityCanvas(QWidget):
    def __init__(self, G, community, parent=None):
        super().__init__(parent)
        self.G = G
        self.community = community
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout(self)
        fig, ax = plt.subplots()
        draw_network(self.G, self.community, ax=ax)
        canvas = FigureCanvas(fig)
        layout.addWidget(canvas)
        self.setLayout(layout)